#include "DTIFiberTrackingCPU.h"

#include <time.h>
#define _USE_MATH_DEFINES
#include <math.h>

#include <teem/ten.h>

#include <string>

#include <omp.h>

CDTIFiberTrackingCPU::CDTIFiberTrackingCPU(void)
{
	m_pDTI = NULL;
	m_pMask = NULL;
	m_pFA = NULL;

	m_width = 0;
	m_height = 0;
	m_depth = 0;

	m_maxLength = 0.0f;
	m_stepSize = 0.0f;
	m_faThreshold = 0.0f;
	m_angleThreshold = 0.0f;

	m_pSamplePoints = NULL;
	m_pSamplePointsWeight = NULL;
}


CDTIFiberTrackingCPU::~CDTIFiberTrackingCPU(void)
{
	SafeDelete(m_pDTI);
	SafeDelete(m_pMask);

	m_width = 0;
	m_height = 0;
	m_depth = 0;

	m_maxLength = 0.0f;
	m_stepSize = 0.0f;
	m_faThreshold = 0.0f;
	m_angleThreshold = 0.0f;
	
	SafeDeleteArray(m_pSamplePoints);
	SafeDeleteArray(m_pSamplePointsWeight);
}

/****************************************************************************************/
/* initalize the fiber tracking class													*/
/* DTIPathname: DTI nrrd file pathname													*/
/* MaskPathname: brain mask nrrd file pathname											*/
/* stepSize: step size used in fiber tracking, may not be used in some tracking method	*/
/* maxLength: maximum fiber length														*/
/* faThreshold: FA threshold to stop the fiber tracking									*/
/* angleThreshold: angle threshold to stop the fiber tracking							*/
/****************************************************************************************/
int CDTIFiberTrackingCPU::InitDTIFiberTracking(const char *DTIPathname, 
											   const char *MaskPathname, 
											   const float stepSize, 
											   const float maxLength,
											   const float faThreshold, 
											   const float angleThreshold)
{
	SafeDelete(m_pDTI);
	SafeDelete(m_pMask);
	
	m_pDTI = new CDTI();
	m_pDTI->OpenDTIFile(DTIPathname, stderr);

	m_pMask = new CScalar();
	m_pMask->OpenScalarFile(MaskPathname, stderr);

	m_stepSize = stepSize;
	m_maxLength = maxLength;
	m_faThreshold = faThreshold;
	m_angleThreshold = angleThreshold;

	m_width = m_pDTI->GetWidth();
	m_height = m_pDTI->GetHeight();
	m_depth = m_pDTI->GetDepth();

	ComputeFA();

	return 0;
}

/************************************************************/
/* Compute FA based on DTI with MLS filter					*/
/************************************************************/
void CDTIFiberTrackingCPU::ComputeFA()
{
	SafeDelete(m_pFA);

	m_pFA = new CScalar();
	m_pFA->CreateScalar(m_width, m_height, m_depth);

	CreateMLSSamplePoints();

	printf("\n=========================\n");
	printf("Computing FA ... 0.00%% done.");
	for (int z = 0; z < (int)m_depth; ++z) {
		omp_set_num_threads(4);
#pragma omp parallel for
		for (int y = 0; y < (int)m_height; ++y) {
			for (int x = 0; x < (int)m_width; ++x) {
				float initDTI[7] = {1, 1, 0, 0, 1, 0, 1};
				float filteredDTI[7];
				CPoint3F pos = CPoint3F((float)x, (float)y, (float)z);
				this->DTIFilter_MLS(&pos, filteredDTI, initDTI);
				float fa = m_pDTI->DTI2FA(filteredDTI);
				m_pFA->SetScalar(x, y, z, fa);
			}
		}

		float ff = (float)(z+1) / (float)(m_depth) * 100.0f;
		printf("\rComputing FA ... %.2f%% done.", ff);
	}
	printf("\rComputing FA ... 100.00%% done.\n");
	printf("=========================\n");
}

/************************************************************/
/* Create MLS sample points									*/
/************************************************************/
void CDTIFiberTrackingCPU::CreateMLSSamplePoints()
{
	SafeDeleteArray(m_pSamplePoints);
	SafeDeleteArray(m_pSamplePointsWeight);
	m_pSamplePoints = new CPoint3F[SAMPLE_POINT_NUM];
	m_pSamplePointsWeight = new float[SAMPLE_POINT_NUM];

	m_pSamplePoints[0].m_x = 0.0f;
	m_pSamplePoints[0].m_y = 0.0f;
	m_pSamplePoints[0].m_z = 0.0f;
	m_pSamplePointsWeight[0] = 1.0f;

	clock_t seed = clock();
	srand(seed);

	for (int i = 1; i < SAMPLE_POINT_NUM; ++i) {
		bool flag = false;
		CPoint3F pt;
		while(!flag) {
			float a = (float)rand() / (float)RAND_MAX * 2.0f * (float)M_PI;
			float b = (float)rand() / (float)RAND_MAX * (float)M_PI;
			float r = (float)rand() / (float)RAND_MAX;

			pt.m_x = r * sinf(a) * sinf(b);
			pt.m_y = r * sinf(a) * cosf(b);
			pt.m_z = r * cosf(a);

			flag = true;
			for (int j = 0; j < i; ++j) {
				float dis = Distance(m_pSamplePoints[j], pt);
				if (dis < 0.25f) {
					flag = false;
					break;
				}
			}
		}

		m_pSamplePoints[i] = pt;
		m_pSamplePointsWeight[i] = InnerProduct(m_pSamplePoints[i], m_pSamplePoints[i]) * 4.0f;
		m_pSamplePointsWeight[i] = expf(m_pSamplePointsWeight[i] / -2.5f);
	}
}


/************************************************************/
/* filter the DTI data by MLS								*/
/************************************************************/
void CDTIFiberTrackingCPU::DTIFilter_MLS(CPoint3F *pos, float *out_DTI, float *in_DTI)
{
	float m[9];
	TEN_T2M(m, in_DTI);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = m_pSamplePoints[i].m_x * m[0] + m_pSamplePoints[i].m_y * m[1] + m_pSamplePoints[i].m_z * m[2];
		samples[i].m_y = m_pSamplePoints[i].m_x * m[3] + m_pSamplePoints[i].m_y * m[4] + m_pSamplePoints[i].m_z * m[5];
		samples[i].m_z = m_pSamplePoints[i].m_x * m[6] + m_pSamplePoints[i].m_y * m[7] + m_pSamplePoints[i].m_z * m[8];
		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = samples[i].m_x / max_length * FILTER_SIZE + pos->m_x;
		samples[i].m_y = samples[i].m_y / max_length * FILTER_SIZE + pos->m_y;
		samples[i].m_z = samples[i].m_z / max_length * FILTER_SIZE + pos->m_z;
	}

	memset(out_DTI, 0, sizeof(float)*7);
	float temp_DTI[7];
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (m_pMask->GetScalar(samples[i]) < 0.5f) {
			continue;
		}
		m_pDTI->GetDTI(samples[i], temp_DTI);

		for (uint j = 0; j < 7; ++j) {
			out_DTI[j] += m_pSamplePointsWeight[i] * temp_DTI[j];
		}
		weight += m_pSamplePointsWeight[i];
	}

	for (uint j = 0; j < 7; ++j) {
		out_DTI[j] = out_DTI[j] / weight;
	}
}

void CDTIFiberTrackingCPU::DTIFilter_TENDMLS(CPoint3F *pos, float *out_DTI, float *in_DTI)
{
	uint xx = uint(pos->m_x + 0.5f);
	uint yy = uint(pos->m_y + 0.5f);
	uint zz = uint(pos->m_z + 0.5f);

	CPoint3F pp((float)xx, (float)yy, (float)zz);
	this->DTIFilter_MLS(&pp, out_DTI, in_DTI);
}


void CDTIFiberTrackingCPU::FiberTracking(FT_METHOD method, 
										 const char *fiberPathname, 
										 const uint scale)
{
	if (FT_FACT == method) {
		FiberTracking_FACT(fiberPathname, scale);
	} else if (FT_STT == method) {
		FiberTracking_STT(fiberPathname, scale);
	} else if (FT_TEND == method) {
		FiberTracking_TEND(fiberPathname, scale);
	} else if (FT_STT_MLS == method) {
		FiberTracking_STTMLS(fiberPathname, scale);
	} else if (FT_TEND_MLS == method) {
		FiberTracking_TENDMLS(fiberPathname, scale);
	} else {
		FiberTracking_STT(fiberPathname, scale);
	}
}


void CDTIFiberTrackingCPU::FiberTracking_FACT(const char *fiberPathname, const uint scale)
{
}

void CDTIFiberTrackingCPU::FiberTracking_STTMLS(const char *fiberPathname, const uint scale)
{
}

void CDTIFiberTrackingCPU::SaveFibers(const CFiber *fibers, const uint w, const uint h, 
									  const char *confPathname, const char *vertPathname)
{
	FILE *confFP = fopen(confPathname, "w");
	FILE *vertFP = fopen(vertPathname, "wb");

	fprintf(confFP, "%d %d\n", w, h);

	for (uint i = 0; i < w*h; ++i) {
		int fnum = fibers[i].m_fCount;
		int bnum = fibers[i].m_bCount;
		CVector3F dir = fibers[i].m_seed.dir;
		fprintf(confFP, "%d\n", 1);
		fprintf(confFP, "%d %d\t", fnum, bnum);
		fprintf(confFP, "%f %f %f\n", dir.m_x, dir.m_y, dir.m_z);

		fwrite(fibers[i].m_pF, sizeof(CPoint3F), fnum, vertFP);
		fwrite(fibers[i].m_pB, sizeof(CPoint3F), bnum, vertFP);
	}

	fclose(confFP);
	fclose(vertFP);
}
